/*
 * AqueousSolution.cpp
 *
 *  Created on: 1 Aug 2011
 *      Author: Allan
 */

#include "AqueousSolution.h"

// C++ includes
#include <fstream>
#include <map>

// Boost includes
#include <boost/foreach.hpp>
#include <boost/algorithm/string.hpp>

// GeoReact includes
#include "Database.h"
#include "Utils.h"

AqueousSolution::AqueousSolution()
{}

AqueousSolution::AqueousSolution(const vector<string>& species) :
species(species)
{
	vector<double> ionicCharges;
	
	// Setting the ionic species
	for(uint ispecies = 0; ispecies < species.size(); ++ispecies)
	{
		// The electrical charge of the current species
		const double z = ElectricalCharge(species[ispecies]);
		
		// Setting the indexes and charges of the ions
		if(z != 0)
		{
			ions.push_back(ispecies);
			ionicCharges.push_back(z);
		}
	}
	
	// Set the index of the water species
	iwater = std::find(species.begin(), species.end(), "H2O") - species.begin();
	
	// Set the charges of the ions (in terms of Eigen vector)
	zi.resize(ions.size()); for(uint ion = 0; ion < ions.size(); ++ion) zi[ion] = ionicCharges[ion];
	
	// Initialize the default activity coefficient models of the aqueous species 
	InitializeActivityModels();
	
	// Initialize the aqueous complexes and the stoichiometric matrix between the complexes and ions
	InitializeComplexes();
}

AqueousSolution::~AqueousSolution()
{}

void AqueousSolution::SetActivityModel(const string& s, const AqueousActivityModel& model)
{
	const unsigned i = std::find(species.begin(), species.end(), s) - species.begin();
	
	activityCoeffFunctions[i] = model;
}

const string AqueousSolution::GetName() const
{
	return "Aqueous";
}

const vector<string> AqueousSolution::GetSpecies() const
{
	return species;
}

const vector<string> AqueousSolution::GetIonicSpecies() const
{
	vector<string> ionicSpecies(ions.size());
	
	for(uint i = 0; i < ions.size(); ++i) ionicSpecies[i] = species[ions[i]];
	
	return ionicSpecies;
}

const VectorXd AqueousSolution::GetIonicCharges() const
{
	return zi;
}

const Index AqueousSolution::GetWaterIndex() const
{
	return iwater;
}

const VectorXd AqueousSolution::Molalities(const VectorXd& n) const
{
	return 55.508/n[iwater] * n;
}

const VectorXd AqueousSolution::MolarFractions(const VectorXd& n) const
{
	return n / n.sum();
}

const VectorXd AqueousSolution::Concentrations(const VectorXd& n) const
{
	return Molalities(n);
}

const VectorXd AqueousSolution::ActivityCoefficients(double T, double P, const VectorXd& n) const
{
	VectorXd g(n.rows());
	
	for(unsigned i = 0; i < (unsigned)g.rows(); ++i) g[i] = activityCoeffFunctions[i](*this, T, P, n);
	
	return g;
}

const VectorXd AqueousSolution::Activities(double T, double P, const VectorXd& n) const
{
	VectorXd a(n.rows());
	
	for(unsigned i = 0; i < (unsigned)a.rows(); ++i) a[i] = n[i] * activityCoeffFunctions[i](*this, T, P, n);
	
	a *= 55.508/n[iwater];
	
	return a;
}

const double AqueousSolution::IonicStrengthEffective(const VectorXd& n) const
{
	double Ie = 0.5 * (zi.array() * zi.array() * IonicMolalities(n).array()).sum();
	
	return Ie;
}

const double AqueousSolution::IonicStrengthStoichiometric(const VectorXd& n) const
{
	double Is = 0.5 * (zi.array() * zi.array() * IonicMolalitiesStoichiometric(n).array()).sum();
	
	return Is;
}

const VectorXd AqueousSolution::IonicMolalities(const VectorXd& n) const
{
	// The ionic molalities
	VectorXd mi(ions.size());
	
	// The molality factor to trasnform number of moles into molality
	const double molalityFactor = 55.508 / n[iwater]; 
	
	// Loop over all the ions
	for(uint i = 0; i < ions.size(); ++i)
		mi[i] = n[ions[i]] * molalityFactor;
	
	return mi;
}

const VectorXd AqueousSolution::IonicMolalitiesStoichiometric(const VectorXd& n) const
{
	// Check if the complexes have been setup
	if(!complexes.empty())
	{
		// The number of ionic and complexes species in the aqueous phase
		const uint numIons = ions.size();
		const uint numComplexes = complexes.size();

		// The molalities of the ions (mi) and complexes (me)
		VectorXd mi(numIons), me(numComplexes);
		
		// The molality factor to trasnform number of moles into molality
		const double molalityFactor = 55.508/n[iwater];
		
		for(uint i = 0; i < numIons; ++i)
			mi[i] = molalityFactor * n[ions[i]];
		
		for(uint e = 0; e < numComplexes; ++e)
			me[e] = molalityFactor * n[complexes[e]];
	
		return mi + vei.transpose() * me;
	}
	else
		return IonicMolalities(n);
}

void AqueousSolution::InitializeActivityModels()
{
	// Set the Setschenow model for uncharged species with coefficient b = 0.1
	activityCoeffFunctions.resize(species.size(), SetschenowModel(0.1));
	
	// Set the ionic Hkf model for charged species
	const vector<string> ionicSpecies = GetIonicSpecies();
	
	for(uint ion = 0; ion < ions.size(); ++ion)
		activityCoeffFunctions[ions[ion]] = HkfIonicModel(ionicSpecies[ion], ionicSpecies);
	
	// Set the water Hkf model for the water species
	activityCoeffFunctions[iwater] = HkfWaterModel(ionicSpecies);
	
	// Set the Drummond model for aqueous CO2 (in case CO2(aq) is present)
	const unsigned iCO2aq = std::find(species.begin(), species.end(), "CO2(aq)") - species.begin();
	
	if(iCO2aq < species.size())
		activityCoeffFunctions[iCO2aq] = DrummondCO2Model();
}

const bool IsComplex(const SpeciesInfo& speciesInfo, const ReactionInfo& reactionInfo)
{
	if(ElectricalCharge(speciesInfo.formula) != 0)
		return false;
	
	for(uint i = 0; i < reactionInfo.products.size(); ++i)
		if(ElectricalCharge(reactionInfo.products[i]) == 0 || reactionInfo.stoichiometries[i] < 0.0)
			return false;
	
	return true;
}

typedef map<string, vector<pair<double, string>>> ComplexMap;

void ReadComplexes(ComplexMap& complexMap)
{
	ifstream file("Databases/complexes.dtb");
	
	string line;
	
	complexMap.clear();
	
	while(!file.eof())
	{
		getline(file, line);
		
		boost::trim_if(line, boost::is_any_of("\n\t "));
		
		vector<string> words = SplitLineAndClean(line, " ,=");
		
		vector<pair<double, string>> composition;
		
		for(unsigned i = 1; i < words.size(); i += 2)
			composition.push_back( {atof(words[i].c_str()), words[i+1]} );
		
		complexMap[words[0]] = composition;
	}
}

void AqueousSolution::InitializeComplexes()
{
	// A complex map object
	ComplexMap complexMap;
	
	// Read all the complexes in the database 'complexes.dtb'
	ReadComplexes(complexMap);
	
	// A function to determine the index of string (s) in the vector of strings (strings)
	auto index = [](const string& s, const vector<string>& strings) { return std::find(strings.begin(), strings.end(), s) - strings.begin(); };
	
	// Loop over all the species in the solution and check which of them is an aqueous complex
	BOOST_FOREACH(auto s, species)
		if(complexMap.find(s) != complexMap.end())
			complexes.push_back( index(s, species) );
	
	// The names of the ionic species
	vector<string> ionicSpecies = GetIonicSpecies();
	
	// Initialise the stoichiometric matrix of the complexes-ions
	vei = MatrixXd::Zero(complexes.size(), ions.size());
	
	// Loop over all the complexes in the aqueous solution
	for(uint c = 0; c < complexes.size(); ++c)
	{
		// The name of the current complex
		const string complex = species[complexes[c]];
		
		typedef pair<double, string> Pair;
		
		// Loop over all pairs (stoichiometry, ion-name)
		BOOST_FOREACH(const Pair& p, complexMap[complex])
		{
			// Get the local index of the ion in the current pair (stoichiometry, ion-name)
			const unsigned i = index(p.second, ionicSpecies);
			
			// Set the entry (c,i) to the stoichiometry of the ion in the current pair (stoichiometry, ion-name)
			vei(c, i) = p.first;
		}
	}
}

